package core

import (
	"bytes"
	"encoding/json"
	"github.com/gorilla/websocket"
	"os/exec"
	"time"
	"watchlog-agent/bean"
	"watchlog-agent/util"
)

type WebSocketClient struct {
	conn         *websocket.Conn
	writeChan    chan []byte
	readChan     chan []byte
	closeChan    chan bool
	done         chan bool
	reconnect    bool
	retrySeconds time.Duration
}

func NewWebSocketClient(urlStr string, reconnect bool, retrySeconds time.Duration) (*WebSocketClient, error) {
	dialer := websocket.DefaultDialer

	conn, _, err := dialer.Dial(urlStr, nil)
	if err != nil {
		return nil, err
	}

	return &WebSocketClient{
		conn:         conn,
		writeChan:    make(chan []byte, 1024),
		readChan:     make(chan []byte, 1024),
		closeChan:    make(chan bool),
		done:         make(chan bool),
		reconnect:    reconnect,
		retrySeconds: retrySeconds,
	}, nil
}

func (c *WebSocketClient) Start() {
	defer func() {
		if c.reconnect {
			go func() {
				time.Sleep(c.retrySeconds * time.Second)
				c.Reconnect()
				c.Start()
			}()
		} else {
			c.done <- true
		}
	}()
	//获取系统mac地址
	mac, _ := util.NewMacAddress()
	go c.read()
	go c.write()
	go c.heartbeat(mac.String())

	<-c.closeChan
	c.conn.Close()
}

func (c *WebSocketClient) Stop() {
	c.closeChan <- true
	<-c.done
	close(c.writeChan)
	close(c.readChan)
	close(c.closeChan)
	close(c.done)
}

func (c *WebSocketClient) Write(data []byte) {
	c.writeChan <- data
}

func (c *WebSocketClient) Read() ([]byte, bool) {
	data, ok := <-c.readChan
	return data, ok
}

func (c *WebSocketClient) Reconnect() {
	dialer := websocket.DefaultDialer
	for {
		util.Log.Info("执行WebSocket重连...")
		conn, _, err := dialer.Dial(c.conn.RemoteAddr().String(), nil)
		if err == nil {
			c.conn = conn
			break
		}
		time.Sleep(c.retrySeconds)
	}
}

func (c *WebSocketClient) read() {
	defer close(c.readChan)
	for {
		_, data, err := c.conn.ReadMessage()
		if err != nil {
			panic(err)
			return
		}
		c.readChan <- data
		var m map[string]interface{}
		if data != nil && len(data) > 0 {
			err := json.Unmarshal(data, &m)
			if err != nil {
				panic(err)
				return
			}
			//执行shell
			shellPath := (m["ShellPath"]).(string)
			//cmd := exec.Command("/bin/sh", "-c", shellPath)
			cmd := exec.Command("/bin/sh", "-c", shellPath)
			var out bytes.Buffer
			var stderr bytes.Buffer
			cmd.Stdout = &out
			cmd.Stderr = &stderr
			err = cmd.Run()
			execResult := bean.ExecResult{
				ShellPath:  shellPath,
				Target:     util.TAG_WEB,
				Result:     false,
				Type:       5,
				ServerName: (m["ServerName"]).(string),
				AgentMac:   (m["AgentMac"]).(string),
			}
			if err != nil {
				execResult.Msg = stderr.String()
				util.Log.Error(err.Error(), stderr.String())
			} else {
				execResult.Result = true
				execResult.Msg = out.String()
				util.Log.Info(out.String())
			}
			toData, _ := json.Marshal(execResult)
			c.Write(toData)
		}
	}
}

func (c *WebSocketClient) write() {
	defer close(c.writeChan)

	for {
		select {
		case data, ok := <-c.writeChan:
			if !ok {
				return
			}

			err := c.conn.WriteMessage(websocket.TextMessage, data)
			if err != nil {
				return
			}
		case <-c.closeChan:
			return
		}
	}
}

func (c *WebSocketClient) heartbeat(mac string) {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()
	ai := bean.Agent{
		Type:     4,
		AgentMac: mac,
	}
	for {
		select {
		case <-ticker.C:
			jsonData, _ := json.Marshal(ai)
			c.Write(jsonData)
		case <-c.closeChan:
			return
		}
	}
}
